use polars::prelude::*;
use calamine::{open_workbook_auto, Reader};
use rust_xlsxwriter::Workbook;
use aes::Aes256;
use cbc::{Decryptor, cipher::{block_padding::Pkcs7, BlockDecryptMut, KeyIvInit}};
use base64::{Engine, engine::general_purpose};

// AES-256-CBC 类型别名
pub type Aes256CbcDec = Decryptor<Aes256>;

/// 解密单个字符串值
fn decrypt_value(encrypted_value: &str, key: &[u8; 32], iv: &[u8; 16]) -> String {
    if let Ok(ciphertext) = general_purpose::STANDARD.decode(encrypted_value) {
        let cipher = Aes256CbcDec::new(key.into(), iv.into());
        let mut buf = ciphertext.clone();
        if let Ok(plaintext) = cipher.decrypt_padded_mut::<Pkcs7>(&mut buf) {
            return String::from_utf8_lossy(plaintext).to_string();
        }
    }
    encrypted_value.to_string()
}

/// 读取 CSV 文件，解密指定列，并导出到新 CSV 文件
/// 
/// # 参数
/// - input_path: 输入 CSV 路径
/// - output_path: 输出 CSV 路径
/// - decrypt_columns: 需解密的列名（如["A", "C"]）
/// - show_columns: 需展示的列名（如["A", "B", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
pub fn process_csv_with_decrypt(
    input_path: &str,
    output_path: &str,
    decrypt_columns: &[&str],
    show_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    // 读取CSV文件到DataFrame
    let mut df = CsvReader::from_path(input_path)?
        .has_header(true)
        .finish()?;
    
    // 对需要解密的列进行处理
    for &col_name in decrypt_columns {
        if df.get_column_names().contains(&col_name) {
            // 获取列数据并解密
            let decrypted_values: Vec<String> = df.column(col_name)?
                .utf8()?
                .into_iter()
                .map(|opt_val| {
                    match opt_val {
                        Some(encrypted_val) => decrypt_value(encrypted_val, key, iv),
                        None => String::new(),
                    }
                })
                .collect();
            
            // 创建新的Series并替换原列
            let new_series = Series::new(col_name, decrypted_values);
            df.with_column(new_series)?;
        }
    }
    
    // 只保留需要展示的列
    let df = df.select(show_columns)?;
    
    // 写入到CSV文件
    let mut file = std::fs::File::create(output_path)?;
    CsvWriter::new(&mut file)
        .include_header(true)
        .finish(&mut df.clone())?;
    
    Ok(())
}

/// 使用calamine读取Excel文件并转换为DataFrame
fn read_excel_to_dataframe(input_path: &str) -> Result<DataFrame, Box<dyn std::error::Error>> {
    let mut workbook = open_workbook_auto(input_path)?;
    let range = workbook.worksheet_range_at(0).ok_or("No sheet found")??;

    // 获取表头
    let headers: Vec<String> = range.rows().next()
        .ok_or("No rows found")?
        .iter()
        .map(|cell| cell.to_string())
        .collect();

    // 收集所有数据行
    let mut columns: Vec<Vec<String>> = vec![Vec::new(); headers.len()];
    
    for row in range.rows().skip(1) {
        for (col_idx, cell) in row.iter().enumerate() {
            if col_idx < columns.len() {
                columns[col_idx].push(cell.to_string());
            }
        }
    }

    // 创建Series并构建DataFrame
    let series: Vec<Series> = headers.iter().zip(columns.iter())
        .map(|(name, data)| Series::new(name, data))
        .collect();

    Ok(DataFrame::new(series)?)
}

/// 将DataFrame写入Excel文件
fn write_dataframe_to_excel(df: &DataFrame, output_path: &str) -> Result<(), Box<dyn std::error::Error>> {
    let mut workbook = Workbook::new();
    let worksheet = workbook.add_worksheet();

    // 写入表头
    let headers = df.get_column_names();
    for (col_idx, header) in headers.iter().enumerate() {
        worksheet.write_string(0, col_idx as u16, *header)?;
    }

    // 写入数据
    let height = df.height();
    for row_idx in 0..height {
        for (col_idx, column) in df.get_columns().iter().enumerate() {
            let value = match column.get(row_idx)? {
                AnyValue::Utf8(s) => s.to_string(),
                AnyValue::Binary(b) => String::from_utf8_lossy(b).to_string(),
                other => format!("{}", other),
            };
            worksheet.write_string((row_idx + 1) as u32, col_idx as u16, &value)?;
        }
    }

    workbook.save(output_path)?;
    Ok(())
}

/// 读取 Excel 文件，解密指定列，并导出到新 Excel 文件
/// 
/// # 参数
/// - input_path: 输入 Excel 路径
/// - output_path: 输出 Excel 路径
/// - decrypt_columns: 需解密的列名（如["A", "C"]）
/// - show_columns: 需展示的列名（如["A", "B", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
pub fn process_excel_with_decrypt(
    input_path: &str,
    output_path: &str,
    decrypt_columns: &[&str],
    show_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    // 检查文件扩展名，决定使用哪种处理方式
    if input_path.ends_with(".csv") {
        return process_csv_with_decrypt(input_path, output_path, decrypt_columns, show_columns, key, iv);
    }
    
    // 读取Excel文件到DataFrame
    let mut df = read_excel_to_dataframe(input_path)?;
    
    // 对需要解密的列进行处理
    for &col_name in decrypt_columns {
        if df.get_column_names().contains(&col_name) {
            // 获取列数据并解密
            let decrypted_values: Vec<String> = df.column(col_name)?
                .utf8()?
                .into_iter()
                .map(|opt_val| {
                    match opt_val {
                        Some(encrypted_val) => decrypt_value(encrypted_val, key, iv),
                        None => String::new(),
                    }
                })
                .collect();
            
            // 创建新的Series并替换原列
            let new_series = Series::new(col_name, decrypted_values);
            df.with_column(new_series)?;
        }
    }
    
    // 只保留需要展示的列
    let df = df.select(show_columns)?;
    
    // 根据输出路径的扩展名决定输出格式
    if output_path.ends_with(".csv") {
        // 输出为CSV
        let mut file = std::fs::File::create(output_path)?;
        CsvWriter::new(&mut file)
            .include_header(true)
            .finish(&mut df.clone())?;
    } else if output_path.ends_with(".xlsx") {
        // 输出为Excel
        write_dataframe_to_excel(&df, output_path)?;
    } else {
        return Err("Unsupported output format. Use .csv or .xlsx extension.".into());
    }
    
    Ok(())
}